package main

type TreeNode struct {
	Val   int
	Left  *TreeNode
	Right *TreeNode
}

var inorderNodes []*TreeNode

func inorder(root *TreeNode) {
	if root == nil {
		return
	}
	inorder(root.Left)
	inorderNodes = append(inorderNodes, root)
	inorder(root.Right)
}

func swap(a, b *TreeNode) {
	tmp := b.Val
	b.Val = a.Val
	a.Val = tmp
}

func recoverTree(root *TreeNode) {
	inorderNodes = nil
	inorder(root)

	var wrongNodes []int
	for i := 0; i < len(inorderNodes)-1; i++ {
		if inorderNodes[i].Val > inorderNodes[i+1].Val {
			wrongNodes = append(wrongNodes, i)
		}
	}
	if len(wrongNodes) == 1 {
		wrongNodes = append(wrongNodes, wrongNodes[0]+1)
	} else if len(wrongNodes) == 2 {
		wrongNodes[1]++
	}
	swap(inorderNodes[wrongNodes[0]], inorderNodes[wrongNodes[1]])
}

func main() {
	root := &TreeNode{
		Val: 1,
		Left: &TreeNode{
			Val: 3,
			Right: &TreeNode{
				Val: 2,
			},
		},
	}
	recoverTree(root)
	println(root)
}
